import torch
from torch import nn
import math
import torch.nn.functional as F
random_torch=torch.rand(4,4)
print(random_torch) 

# class TokenEmbedding(nn.model):
#     #词汇表大小，embedding的维度
#     def __init__():
#         super()
#     #将输入的词汇表索引转换为指定维度的Embedding
class PositionEmbedding(nn.Module):
    def __init__(self,d_model,max_len,device):
        #直接调用父类函数的__init__()
        super(PositionEmbedding,self).__init__()
        self.encoding=torch.zeros(max_len,d_model,device=device)
        #不计算梯度
        self.encoding.requires_grad=False
        pos=torch.arange(0,max_len,device=device)
        print("pos=",pos)
        pos=pos.float().unsqueeze(dim=1)
#嵌入维度
d_model=512#词向量维度，比如I → [0.1, 0.2, ..., 0.9]
max_len=100
n_head=2 #多头注意力机制的头数
device=torch.device("cuda"if torch.cuda.is_available() else "cpu")
pos_embed=PositionEmbedding(d_model,max_len,device)

x=torch.rand(128,32,512)

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model,n_head):
        super().__init__()
        self.n_head=n_head
        self.d_model=d_model
        self.w_q=nn.Linear(d_model,d_model)
        print("w_q=",self.w_q)
        self.w_k=nn.Linear(d_model,d_model)
        self.w_v=nn.Linear(d_model,d_model)
        self.w_combine=nn.Linear(d_model,d_model)
        self.softmax=nn.Softmax(dim=1)
        self.dropout = nn.Dropout(0.1)

    def forward(self,q,k,v,mask=None):
        batch_size,seq_len,dismesion=q.shape
        # 1. 线性变换
        #拆分的头数为512//2个头
        head_dim=self.d_model//self.n_head
        Q = self.w_q(q)  # (batch, seq_len, d_model)
        K = self.w_k(k)  # (batch, seq_len, d_model)
        V = self.w_v(v)  # (batch, seq_len, d_model)
        # 2. 切分为多头
        #表示：batch_size个批量，3 个词，每个词拆成 2 个头，每个头拥有256(head_dim)维的词向量。
        #transpose(1, 2)将seq_len和n_head维度交换位置,变为2个词和3个头
        #transpose 是为了把“每个词有多个头”的结构，变成“每个头看整句话”的结构
        #从而让多个 attention head 并行计算，是高效实现多头注意力的标准做法。
        Q = Q.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)  # (batch, n_head, seq_len, head_dim)
        K = K.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)
        V = V.view(batch_size, seq_len, self.n_head, head_dim).transpose(1, 2)
        scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(head_dim)
        # 4. 掩码（可选）
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        # 5. Softmax和Dropout
        attn_weights = self.softmax(scores)
        attn_weights = self.dropout(attn_weights)
        # 6. 加权求和
        output = torch.matmul(attn_weights, V)  # (batch, n_head, seq_len, head_dim)
        # 7. 合并多头
        output = output.transpose(1, 2).contiguous()  # (batch, seq_len, n_head, head_dim)
        output = output.view(batch_size, seq_len, self.d_model)  # (batch, seq_len, d_model)
        # 8. 最终线性变换
        output = self.w_combine(output)
        
        return output
        
class LayerNorm(nn.Module):
    def __init__(self,d_model, eps=1e-6):#数值稳定性
        super().__init__()
        self.gamma=nn.Parameter(torch.ones(d_model))  # 可学习参数 nn.Parameter 是 PyTorch 中的一个核心类，
        #​用于将张量（Tensor）标记为模型的可学习参数。它的主要作用是告诉 PyTorch 的优化器（如 Adam、SGD）在训练过程中需要更新这些张量的值（即通过梯度下降进行优化）。
        self.beta=nn.Parameter(torch.zeros(d_model))
        self.eps = eps
    def forward(self, x):
        mean=x.mean(dim=-1, keepdim=True)  # 计算均值维度不变
        var=x.var(dim=-1,keepdim=True)  # 计算方差维度不变
        out=(x-mean)/torch.sqrt(var+self.eps)  # 标准化
        out=self.gamma*out+self.beta
        return out
    

d_model = 4  # 特征维度需与输入数据的最后一维一致
layernorm = LayerNorm(d_model=d_model, eps=1e-6)

# 打印可学习参数 gamma 和 beta
print("gamma:", layernorm.gamma)  # 初始值为全1
print("beta:", layernorm.beta)  





LN=LayerNorm()
LN.forward
d_model = 512
n_head = 8
mha = MultiHeadAttention(d_model, n_head)   